package io.github.kevinzhwl.edgetts.util;

import org.apache.commons.lang3.StringUtils;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import java.io.FileOutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;

public class TTSWssClient extends WebSocketClient {

  private String audioFilePath;

  public TTSWssClient(String serverUri, Map<String, String> httpHeaders, String audioFilePath) throws URISyntaxException {
    super(new URI(serverUri), httpHeaders);
    this.audioFilePath = audioFilePath;
  }

  @Override
  public void onOpen(ServerHandshake handshakedata) {

  }

  @Override
  public void onMessage(String message) {
    if (message.contains("Path:turn.end")) {
      close();
    }
  }

  @Override
  public void onMessage(ByteBuffer originBytes) {
    saveByContentType(originBytes);
  }


  /**
   * This method directly specifies the file header marker, which makes it faster. However, if the format changes, it may become unusable.
   * See: https://github.com/microsoft/cognitive-services-speech-sdk-js/blob/d071d11/src/common.speech/WebsocketMessageFormatter.ts#L46
   *
   * @param originBytes
   */
  public void saveByContentType(ByteBuffer originBytes) {
    byte[] rawOrigin = originBytes.array();
    byte[] lengthBytesHead = Arrays.copyOfRange(rawOrigin, 0, 2);
    final int dataLength = toIntBig(lengthBytesHead, 2);

    byte[] voiceBytesHead = Arrays.copyOfRange(rawOrigin, 2, 2 + dataLength);
    final String str = StringUtils.toEncodedString(voiceBytesHead, StandardCharsets.UTF_8);

    final int dataOffset;
    if (str.contains("Content-Type")) {
      dataOffset = dataLength + 2;
    } else {
      dataOffset = 105;
    }

    byte[] voiceBytesRemoveHead = Arrays.copyOfRange(rawOrigin, dataOffset, rawOrigin.length);
    try (FileOutputStream fos = new FileOutputStream(audioFilePath, true)) {
      fos.write(voiceBytesRemoveHead);
      fos.flush();
    } catch (Exception ex) {
      ex.printStackTrace();
    }
  }

  public int toIntBig(byte[] b, int bLength) {
    int rv = 0;
    for (int i = 0; i < bLength; i++) {
      rv <<= 8;
      rv += (b[i] & 0xff);
    }
    return rv;
  }

  @Override
  public void onClose(int code, String reason, boolean remote) {
  }

  @Override
  public void onError(Exception ex) {
    ex.printStackTrace();
  }
}
